from PIL import Image, ImageDraw
import os
import random
import math

def ensure_dir(directory):
    """确保目录存在"""
    if not os.path.exists(directory):
        os.makedirs(directory)

def generate_arrows():
    """生成箭头图形 - 用于方向识别训练"""
    print("开始生成箭头图形...")
    
    ensure_dir('../arrows')
    
    directions = ['up', 'down', 'left', 'right']
    colors = ['red', 'blue', 'yellow', 'green']
    color_values = {
        'red': (255, 0, 0),
        'blue': (0, 0, 255),
        'yellow': (255, 255, 0),
        'green': (0, 255, 0)
    }
    
    arrow_count = 0
    
    # 每个方向生成500张，总共2000张
    for direction in directions:
        for color_name in colors:
            print(f"生成 {color_name} {direction} 箭头...")
            for i in range(125):  # 每个颜色方向组合生成125个，4颜色×4方向×125=2000
                # 创建白色背景图片
                img = Image.new('RGB', (224, 224), (255, 255, 255))
                draw = ImageDraw.Draw(img)
                
                color = color_values[color_name]
                
                # 添加随机变化
                offset_x = random.randint(-15, 15)
                offset_y = random.randint(-15, 15)
                size_scale = random.uniform(0.8, 1.2)
                
                # 根据方向绘制不同的箭头
                base_size = 30 * size_scale
                if direction == 'up':
                    # 向上箭头
                    center_x, center_y = 112 + offset_x, 112 + offset_y
                    points = [
                        (center_x, center_y - base_size),
                        (center_x - base_size, center_y + base_size/2),
                        (center_x - base_size/3, center_y + base_size/2),
                        (center_x - base_size/3, center_y + base_size),
                        (center_x + base_size/3, center_y + base_size),
                        (center_x + base_size/3, center_y + base_size/2),
                        (center_x + base_size, center_y + base_size/2)
                    ]
                elif direction == 'down':
                    # 向下箭头
                    center_x, center_y = 112 + offset_x, 112 + offset_y
                    points = [
                        (center_x, center_y + base_size),
                        (center_x - base_size, center_y - base_size/2),
                        (center_x - base_size/3, center_y - base_size/2),
                        (center_x - base_size/3, center_y - base_size),
                        (center_x + base_size/3, center_y - base_size),
                        (center_x + base_size/3, center_y - base_size/2),
                        (center_x + base_size, center_y - base_size/2)
                    ]
                elif direction == 'left':
                    # 向左箭头
                    center_x, center_y = 112 + offset_x, 112 + offset_y
                    points = [
                        (center_x - base_size, center_y),
                        (center_x + base_size/2, center_y - base_size),
                        (center_x + base_size/2, center_y - base_size/3),
                        (center_x + base_size, center_y - base_size/3),
                        (center_x + base_size, center_y + base_size/3),
                        (center_x + base_size/2, center_y + base_size/3),
                        (center_x + base_size/2, center_y + base_size)
                    ]
                elif direction == 'right':
                    # 向右箭头
                    center_x, center_y = 112 + offset_x, 112 + offset_y
                    points = [
                        (center_x + base_size, center_y),
                        (center_x - base_size/2, center_y - base_size),
                        (center_x - base_size/2, center_y - base_size/3),
                        (center_x - base_size, center_y - base_size/3),
                        (center_x - base_size, center_y + base_size/3),
                        (center_x - base_size/2, center_y + base_size/3),
                        (center_x - base_size/2, center_y + base_size)
                    ]
                
                draw.polygon(points, fill=color)
                
                # 保存图片
                filename = f'../arrows/{direction}_{color_name}_{i:03d}.png'
                img.save(filename)
                arrow_count += 1
                
                # 每生成100张显示进度
                if arrow_count % 100 == 0:
                    print(f"  已生成 {arrow_count}/2000 张箭头图片")
    
    print(f"箭头图形生成完成！共生成 {arrow_count} 个箭头")

def generate_variations():
    """生成一些变体箭头，增加数据多样性"""
    print("开始生成变体箭头图片...")
    
    ensure_dir('../variations')
    
    directions = ['up', 'down', 'left', 'right']
    colors = ['red', 'blue', 'yellow', 'green']
    color_values = {
        'red': (255, 0, 0),
        'blue': (0, 0, 255),
        'yellow': (255, 255, 0),
        'green': (0, 255, 0)
    }
    
    # 生成400张变体（总共2400张）
    variation_count = 0
    for i in range(100):  # 每个方向25张变体
        for direction in directions:
            for color_name in colors:
                if variation_count >= 400:  # 总共400张变体
                    break
                    
                # 创建随机背景色（浅色）
                bg_color = (
                    random.randint(200, 255),
                    random.randint(200, 255), 
                    random.randint(200, 255)
                )
                
                img = Image.new('RGB', (224, 224), bg_color)
                draw = ImageDraw.Draw(img)
                
                color = color_values[color_name]
                
                # 更大的随机变化
                offset_x = random.randint(-20, 20)
                offset_y = random.randint(-20, 20)
                size_scale = random.uniform(0.7, 1.3)
                rotation = random.randint(-10, 10)  # 轻微旋转
                
                base_size = 30 * size_scale
                if direction == 'up':
                    points = get_rotated_points([
                        (112, 50), (80, 120), (95, 120), (95, 174), (129, 174), (129, 120), (144, 120)
                    ], rotation, offset_x, offset_y)
                elif direction == 'down':
                    points = get_rotated_points([
                        (112, 174), (80, 104), (95, 104), (95, 50), (129, 50), (129, 104), (144, 104)
                    ], rotation, offset_x, offset_y)
                elif direction == 'left':
                    points = get_rotated_points([
                        (50, 112), (120, 80), (120, 95), (174, 95), (174, 129), (120, 129), (120, 144)
                    ], rotation, offset_x, offset_y)
                elif direction == 'right':
                    points = get_rotated_points([
                        (174, 112), (104, 80), (104, 95), (50, 95), (50, 129), (104, 129), (104, 144)
                    ], rotation, offset_x, offset_y)
                
                draw.polygon(points, fill=color)
                
                filename = f'../variations/arrow_{direction}_{color_name}_{i:02d}.png'
                img.save(filename)
                variation_count += 1
    
    print(f"变体箭头生成完成！共生成 {variation_count} 个变体箭头")

def get_rotated_points(points, angle, offset_x, offset_y):
    """对点集进行旋转和平移"""
    import math
    angle_rad = math.radians(angle)
    cos_angle = math.cos(angle_rad)
    sin_angle = math.sin(angle_rad)
    
    rotated_points = []
    center_x, center_y = 112, 112
    
    for x, y in points:
        # 平移至中心
        x_centered = x - center_x
        y_centered = y - center_y
        
        # 旋转
        x_rotated = x_centered * cos_angle - y_centered * sin_angle
        y_rotated = x_centered * sin_angle + y_centered * cos_angle
        
        # 平移回原位置并添加偏移
        x_final = x_rotated + center_x + offset_x
        y_final = y_rotated + center_y + offset_y
        
        rotated_points.append((x_final, y_final))
    
    return rotated_points

def main():
    """主函数"""
    print("=" * 50)
    print("K210 箭头图片生成器")
    print("=" * 50)
    
    # 生成箭头图片
    generate_arrows()
    print("-" * 30)
    
    # 生成变体箭头
    generate_variations()
    print("-" * 30)
    
    print("所有箭头图片生成完成！")
    print("图片保存在以下文件夹：")
    print("  - arrows/: 2000张基础箭头")
    print("  - variations/: 400张变体箭头")

if __name__ == "__main__":
    main()